{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FYLyuStTYesc"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "PVm-iEoxYesf"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3MPf91rVYesq"
      },
      "source": [
        "# TensorFlow graph optimization with Grappler"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zmNCsZlgYesr"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/graph_optimization\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/graph_optimization.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/guide/graph_optimization.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/graph_optimization.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l0qacLgyYess"
      },
      "source": [
        "## Overview\n",
        "\n",
        "TensorFlow uses both graph and eager executions to execute computations. A `tf.Graph` contains a set of `tf.Operation` objects (ops) which represent units of computation and `tf.Tensor` objects which represent the units of data that flow between ops.\n",
        "\n",
        "Grappler is the default graph optimization system in the TensorFlow runtime. Grappler applies optimizations in graph mode (within `tf.function`) to improve the performance of your TensorFlow computations through graph simplifications and other high-level optimizations such as inlining function bodies to enable inter-procedural optimizations. Optimizing the `tf.Graph` also reduces the device peak memory usage and improves hardware utilization by optimizing the mapping of graph nodes to compute resources. \n",
        "\n",
        "Use `tf.config.optimizer.set_experimental_options()` for finer control over your `tf.Graph` optimizations.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A-zkJgR5Yesw"
      },
      "source": [
        "## Available graph optimizers\n",
        "\n",
        "Grappler performs graph optimizations through a top-level driver called the `MetaOptimizer`. The following graph optimizers are available with TensorFlow:  \n",
        "\n",
        "* *Constant folding optimizer -* Statically infers the value of tensors when possible by folding constant nodes in the graph and materializes the result using constants.\n",
        "* *Arithmetic optimizer -* Simplifies arithmetic operations by eliminating common subexpressions and simplifying arithmetic statements.  \n",
        "* *Layout optimizer -* Optimizes tensor layouts to execute data format dependent operations such as convolutions more efficiently.\n",
        "* *Remapper optimizer -* Remaps subgraphs onto more efficient implementations by replacing commonly occurring subgraphs with optimized fused monolithic kernels.\n",
        "* *Memory optimizer -* Analyzes the graph to inspect the peak memory usage for each operation and inserts CPU-GPU memory copy operations for swapping GPU memory to CPU to reduce the peak memory usage.\n",
        "* *Dependency optimizer -* Removes or rearranges control dependencies to shorten the critical path for a model step or enables other\n",
        "optimizations. Also removes nodes that are effectively no-ops such as Identity.\n",
        "* *Pruning optimizer -* Prunes nodes that have no effect on the output from the graph. It is usually run first to reduce the size of the graph and speed up processing in other Grappler passes.\n",
        "* *Function optimizer -* Optimizes the function library of a TensorFlow program and inlines function bodies to enable other inter-procedural optimizations.\n",
        "* *Shape optimizer -* Optimizes subgraphs that operate on shape and shape related information.\n",
        "* *Autoparallel optimizer -* Automatically parallelizes graphs by splitting along the batch dimension. This optimizer is turned OFF by default.\n",
        "* *Loop optimizer -* Optimizes the graph control flow by hoisting loop-invariant subgraphs out of loops and by removing redundant stack operations in loops. Also optimizes loops with statically known trip counts and removes statically known dead branches in conditionals.\n",
        "* *Scoped allocator optimizer -* Introduces scoped allocators to reduce data movement and to consolidate some operations.\n",
        "* *Pin to host optimizer -* Swaps small operations onto the CPU. This optimizer is turned OFF by default. \n",
        "* *Auto mixed precision optimizer -* Converts data types to float16 where applicable to improve performance. Currently applies to GPUs and the latest Intel Xeon CPUs.\n",
        "* *Debug stripper -* Strips nodes related to debugging operations such as `tf.debugging.Assert`, `tf.debugging.check_numerics`, and `tf.print` from the graph. This optimizer is turned OFF by default."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WZAUsxyWYess"
      },
      "source": [
        "## Setup\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6BRIDzO6ypoY"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import timeit\n",
        "import traceback\n",
        "import contextlib\n",
        "\n",
        "\n",
        "import tensorflow as tf"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1O-XL1nxJX0X"
      },
      "source": [
        "Create a context manager to easily toggle optimizer states."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uRuhVoAlYesz"
      },
      "outputs": [],
      "source": [
        "@contextlib.contextmanager\n",
        "def options(options):\n",
        "  old_opts = tf.config.optimizer.get_experimental_options()\n",
        "  tf.config.optimizer.set_experimental_options(options)\n",
        "  try:\n",
        "    yield\n",
        "  finally:\n",
        "    tf.config.optimizer.set_experimental_options(old_opts)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E2o4kZtK0DoA"
      },
      "source": [
        "## Compare execution performance with and without Grappler\n",
        "\n",
        "TensorFlow 2 and beyond executes eagerly by default. Use `tf.function` to switch the default execution to Graph mode. Grappler runs automatically in the background to apply the graph optimizations above and improve execution performance. \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3sh8RoLJ96IT"
      },
      "source": [
        "### Constant folding optimizer\n",
        "\n",
        "As a preliminary example, consider a function which performs operations on constants and returns an output."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jOW_OSzMJEvN"
      },
      "outputs": [],
      "source": [
        "def test_function_1():\n",
        "  @tf.function\n",
        "  def simple_function(input_arg):\n",
        "    print('Tracing!')\n",
        "    a = tf.constant(np.random.randn(2000,2000), dtype = tf.float32)\n",
        "    c = a\n",
        "    for n in range(50):\n",
        "      c = c@a\n",
        "    return tf.reduce_mean(c+input_arg)\n",
        "\n",
        "  return simple_function"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tFVgUhhzLKIo"
      },
      "source": [
        "Turn off the constant folding optimizer and execute the function:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KDMGsOHrJqKD"
      },
      "outputs": [],
      "source": [
        "with options({'constant_folding': False}):\n",
        "  print(tf.config.optimizer.get_experimental_options())\n",
        "  simple_function = test_function_1()\n",
        "  # Trace once\n",
        "  x = tf.constant(2.2)\n",
        "  simple_function(x)\n",
        "  print(\"Vanilla execution:\", timeit.timeit(lambda: simple_function(x), number = 1), \"s\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ykMXfo8qO41z"
      },
      "source": [
        "Enable the constant folding optimizer and execute the function again to observe a speed-up in function execution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "73pc0gfFKY8l"
      },
      "outputs": [],
      "source": [
        "with options({'constant_folding': True}):\n",
        "  print(tf.config.optimizer.get_experimental_options())\n",
        "  simple_function = test_function_1()\n",
        "  # Trace once\n",
        "  x = tf.constant(2.2)\n",
        "  simple_function(x)\n",
        "  print(\"Constant folded execution:\", timeit.timeit(lambda: simple_function(x), number = 1), \"s\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "83w8rfcRVhWb"
      },
      "source": [
        "### Debug stripper optimizer\n",
        "\n",
        "Consider a simple function that checks the numeric value of its input argument and returns it. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j2DvqEr8haut"
      },
      "outputs": [],
      "source": [
        "def test_function_2():\n",
        "  @tf.function\n",
        "  def simple_func(input_arg):\n",
        "    output = input_arg\n",
        "    tf.debugging.check_numerics(output, \"Bad!\")\n",
        "    return output\n",
        "  return simple_func"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ywKG3WRbpYB8"
      },
      "source": [
        "First, execute the function with the debug stripper optimizer turned off. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LsE-y6iQWSwH"
      },
      "outputs": [],
      "source": [
        "test_func = test_function_2()\n",
        "p1 = tf.constant(float('inf'))\n",
        "try:\n",
        "  test_func(p1)\n",
        "except tf.errors.InvalidArgumentError as e:\n",
        "  traceback.print_exc(limit=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "URHpboM8xLN6"
      },
      "source": [
        "`tf.debugging.check_numerics` raises an invalid argument error because of the `Inf` argument to `test_func`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CuPSha9YmJRo"
      },
      "source": [
        "Enable the debug stripper optimizer and execute the function again. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UPJ7ygHnWP6B"
      },
      "outputs": [],
      "source": [
        "with options({'debug_stripper': True}):\n",
        "  test_func2 = test_function_2()\n",
        "  p1 = tf.constant(float('inf'))\n",
        "  try:\n",
        "    test_func2(p1)\n",
        "  except tf.errors.InvalidArgumentError as e:\n",
        "    traceback.print_exc(limit=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nAsESNUB1QpI"
      },
      "source": [
        "The debug stripper optimizer strips the `tf.debug.check_numerics` node from the graph and executes the function without raising any errors. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wHC6tR9GvFgW"
      },
      "source": [
        "## Summary\n",
        "\n",
        "The TensorFlow runtime uses Grappler to optimize graphs automatically before execution. Use `tf.config.optimizer.set_experimental_options` to enable or disable the various graph optimizers. \n",
        "\n",
        "For more information on Grappler, see <a href=\"http://web.stanford.edu/class/cs245/slides/TFGraphOptimizationsStanford.pdf\" class=\"external\">TensorFlow Graph Optimizations</a>."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "graph_optimization.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
